from user_utils import Tools
import gradio as gr
from tqdm import tqdm

tools = Tools()

# 全局属性
# 初始化状态值为字典以保存用户选择的数据库和表
state = gr.State(value={'db': None, 'tables': None, 'mode': 'sql'})

# 显示所有的数据库


def show_databases():
    return tools.show_databases()

# 选择数据库触发的操作


def database_dropdown_change(db_name):
    """当用户选择数据库时触发的函数，显示数据库中的表"""
    print(f"\n=== database_dropdown_change 开始执行 ===\n选择的数据库: '{db_name}'")
    try:
        # 检查数据库名称是否有效
        if not db_name or db_name == "":
            print("  错误: 数据库名称为空")
            gr.Warning("未选择数据库")
            return gr.CheckboxGroup(
                choices=["请选择一个有效的数据库"], 
                label='选择查询表', 
                visible=True, 
                interactive=False
            )
        
        # 全局保存选中的数据库
        state.value['db'] = db_name
        print(f"  状态已更新: state.value['db'] = {db_name}")
        
        # 着重断点: 查询数据库下所有的数据表
        print(f"  开始查询数据库 {db_name} 中的表")
        tables = tools.show_tables(db_name)
        print(f"  获取到的表: {tables}")
        
        # 特别检查是否有数据表
        if tables is None:
            print("  错误: 数据表查询返回None")
            gr.Warning(f"数据库 {db_name} 查询表时出错")
            return gr.CheckboxGroup(
                choices=["查询数据表出错"], 
                label='选择查询表 (出错)', 
                visible=True, 
                interactive=False
            )
        
        # 检查表格列表是否为空
        if len(tables) == 0:
            print(f"  警告: 数据库 {db_name} 中没有找到任何表")
            gr.Warning(f"数据库 {db_name} 没有找到任何表")
            return gr.CheckboxGroup(
                choices=["此数据库没有表"], 
                label=f'选择查询表 (数据库 {db_name} 没有表)', 
                visible=True, 
                interactive=False
            )
        
        # 成功获取到表格
        print(f"  成功: 找到 {len(tables)} 个表: {tables}")
        gr.Info(f"成功！找到 {len(tables)} 个表")
        
        # 返回更新后的CheckboxGroup组件
        result = gr.CheckboxGroup(
            choices=tables, 
            label=f'选择查询表 (共{len(tables)}个表)', 
            visible=True, 
            interactive=True
        )
        print(f"=== database_dropdown_change 执行完成 ===\n")
        return result
        
    except Exception as e:
        print(f"  严重错误: 获取表格时出错: {e}")
        gr.Error(f"错误: {str(e)}")
        # 出错时也要返回一个可见的组件，显示错误信息
        return gr.CheckboxGroup(
            choices=[f"错误: {str(e)}"], 
            label='选择查询表 (出错)', 
            visible=True, 
            interactive=False
        )

# 查询数据表改变后触发的操作


def query_tables_change(select_tables):
    print(f"\n=== 表格选择变化 ===\n选中的表: {select_tables}")
    
    # 检查当前数据库是否选择
    if 'db' not in state.value or not state.value['db']:
        print("  错误: 当前没有选择数据库, 无法保存表格选择")
        gr.Warning("请先选择数据库再选择表格")
        state.value['tables'] = []
        return
        
    # 检查选择的表格是否有效
    if select_tables and len(select_tables) > 0:
        # 检查选择的第一个项目是否是错误消息
        if select_tables[0].startswith('错误:') or select_tables[0] == '此数据库没有表' or select_tables[0] == '请先选择数据库':
            print(f"  警告: 选中的是错误消息而非有效表名: {select_tables[0]}")
            state.value['tables'] = []
            return
            
        # 保存选中的数据表
        state.value['tables'] = select_tables  # 保存为列表而不是元组
        print(f"  已保存 {len(select_tables)} 个选中的表: {select_tables}")
        
        # 同步表格选择到app.py中
        try:
            from app import app_state
            app_state['tables'] = select_tables
            print(f"  已同步表格选择到app模块: {len(select_tables)}个表")
        except Exception as e:
            print(f"  警告: 无法同步表格选择到app模块: {e}")
        
        if len(select_tables) > 0:
            gr.Info(f"已选择 {len(select_tables)} 个表")
    else:
        # 清除选中的表
        state.value['tables'] = []
        print("  清除选中的表")
        
        # 同步清除操作到app.py
        try:
            from app import app_state
            app_state['tables'] = []
            print("  已同步清除表格选择到app模块")
        except Exception as e:
            print(f"  警告: 无法同步清除表格选择到app模块: {e}")

# 生成模式改变后触发的操作


def generate_model_change(option):
    if option == 'SQL':
        state.value['mode'] = 'sql'
        gr.Info('SQL Mode')
        return gr.Button(value='Generate SQL', visible=True, scale=1, variant='primary'), gr.Button(value='Generate Code', visible=False)
    elif option == 'Code':
        state.value['mode'] = 'code'
        gr.Info('Code Mode')
        return gr.Button(value='Generate SQL', visible=False), gr.Button(value='Generate Code', visible=True, scale=1, variant='secondary')

# 点击生成sql按钮触发的操作


def generate_sql_bt_click(prompt):
    print(f"\n=== 生成SQL按钮点击 ===\n用户提示: '{prompt}'")
    
    # 输出当前state的完整状态信息
    print(f"  当前state状态: {state.value}")
    
    # 检查用户提示
    if not prompt or prompt.strip() == "":
        print("  错误: 用户提示为空")
        gr.Warning("请输入查询提示")
        return gr.Code(label='Query SQL', 
                       value='-- 请输入查询问题', 
                       language='sql-mySQL', 
                       interactive=True, 
                       visible=True), gr.Dataframe(label='Query Result', value=None, visible=True)
    
    # 获取使用的数据库 - 强化验证
    db_name = state.value.get('db')
    print(f"  当前数据库: '{db_name}'")
    
    # 如果状态中没有数据库，尝试从全局状态获取
    if not db_name:
        try:
            from app import app_state
            db_name = app_state.get('db')
            if db_name:
                print(f"  从全局状态恢复数据库选择: {db_name}")
                state.value['db'] = db_name  # 更新当前状态
        except Exception as e:
            print(f"  无法从全局状态获取数据库: {e}")
    
    if not db_name:
        print("  错误: 没有选择数据库")
        gr.Warning("请先选择数据库")
        return gr.Code(label='Query SQL', 
                       value='-- 错误: 请先选择数据库', 
                       language='sql-mySQL', 
                       interactive=True, 
                       visible=True), gr.Dataframe(label='Query Result', value=None, visible=True)
    
    # 获取选中的表 - 强化验证
    tables = state.value.get('tables')
    
    # 如果本地状态中没有表格信息，尝试从全局状态获取
    if not tables or len(tables) == 0:
        try:
            from app import app_state
            tables_from_app = app_state.get('tables')
            if tables_from_app and len(tables_from_app) > 0:
                tables = tables_from_app
                state.value['tables'] = tables  # 更新当前状态
                print(f"  从全局状态恢复表格选择: {tables}")
        except Exception as e:
            print(f"  无法从全局状态获取表格: {e}")
            
    print(f"  选中的表: {tables}")
    
    if not tables or len(tables) == 0:
        print("  错误: 没有选择表")
        gr.Warning("请选择至少一个表格")
        return gr.Code(label='Query SQL', 
                       value='-- 错误: 请选择至少一个表格', 
                       language='sql-mySQL', 
                       interactive=True, 
                       visible=True), gr.Dataframe(label='Query Result', value=None, visible=True)
    
    # 输出当前已验证的状态
    print(f"  已验证的数据库: {db_name} | 表格: {tables}")
    
    try:
        # 获取表的描述信息
        tables_info = ''
        for table in tqdm(tables, total=len(tables), desc='Get Table Info'):
            table_info = tools.get_table_describe(db_name, table)
            if table_info:
                tables_info += f'\n{table_info}'
        
        if not tables_info.strip():
            print("  错误: 无法获取表的描述信息")
            gr.Warning("获取表结构信息失败")
            return gr.Code(label='Query SQL', 
                           value='-- 错误: 无法获取表结构信息', 
                           language='sql-mySQL', 
                           interactive=True, 
                           visible=True), gr.Dataframe(label='Query Result', value=None, visible=True)

        # 生成SQL
        full_prompt = tables_info.strip() + "\n\n" + prompt
        print(f"  向AI发送提示 (前100字符): {full_prompt[:100]}...")
        
        response = tools.generate_sql(full_prompt)
        print("  AI响应已返回")

        # 从 markdown 代码块中提取出 sql 语句
        sql = tools.extract_sql_statements(response)
        if not sql or sql.strip() == "":
            print("  错误: 无法从响应中提取SQL")
            gr.Warning("AI生成的响应中没有发现有效的SQL")
            return gr.Code(label='Query SQL', 
                           value=f'-- AI响应中未找到SQL\n-- 原始响应: {response}', 
                           language='sql-mySQL', 
                           interactive=True, 
                           visible=True), gr.Dataframe(label='Query Result', value=None, visible=True)
        
        print(f"  提取的SQL: {sql[:100]}...")
        
        # 执行 SQL 查询
        print(f"  将在数据库 '{db_name}' 中执行SQL")
        query_result = tools.execute_sql(db_name, sql)
        print("  SQL执行成功")

        return gr.Code(label='Query SQL', 
                       value=sql, 
                       language='sql-mySQL', 
                       interactive=True, 
                       visible=True), gr.Dataframe(label='Query Result', 
                                                  value=query_result, 
                                                  visible=True)
                                                  
    except Exception as e:
        print(f"  错误: 生成或执行SQL时出错: {e}")
        gr.Error(f"错误: {str(e)}")
        return gr.Code(label='Query SQL', 
                       value=f'-- 错误: {str(e)}', 
                       language='sql-mySQL', 
                       interactive=True, 
                       visible=True), gr.Dataframe(label='Query Result', value=None, visible=True)

# 点击生成code按钮触发的操作


def generate_code_bt_click(user_prompt, df):
    if user_prompt:
        # 查询df的信息
        df_info = str(df)

        prompt = f'Pandas中的DataFrame数据的info信息为:{df_info},从中提取出关键信息,使用Python中的matplotlib实现数据可视化,生成的代码最后只返回plt对象,具体要求为:{user_prompt}'

        code = tools.extract_code_statements(tools.generate_code(prompt))

        return gr.Code(label='Python Code', interactive=True, visible=True, scale=1, language='python', value=code), gr.Button(value='Execute Python Code', visible=True)

# 点击执行code按钮触发的操作


def execute_code_bt_click(code):
    import os
    import traceback
    import matplotlib.pyplot as plt
    
    try:
        # 确保可视化目录存在
        vis_dir = 'visualization'
        if not os.path.exists(vis_dir):
            os.makedirs(vis_dir)
            print(f"  创建可视化目录: {vis_dir}")
        
        # 执行前关闭所有可能的图表
        plt.close('all')
        
        if not code or code.strip() == "":
            print("  错误: 没有代码要执行")
            gr.Warning("没有代码要执行")
            
            # 创建错误提示图
            plt.figure(figsize=(6, 4))
            plt.text(0.5, 0.5, '没有代码要执行', ha='center', va='center', fontsize=18)
            plt.axis('off')
            error_file = os.path.join(vis_dir, 'no_code_error.png')
            plt.savefig(error_file)
            plt.close()
            
            return gr.Gallery(label='Visualization Error', 
                              value=[error_file], 
                              preview=True, 
                              visible=True, 
                              columns=1, 
                              rows=1, 
                              object_fit='contain', 
                              height='auto')
        
        print(f"  开始执行代码 (代码长度: {len(code)}):")
        print(f"  代码前50个字符: {code[:50]}...")
        
        # 使用全局变量执行代码，聚焦在plt对象
        exec_globals = {}
        exec_globals.update(globals())
        exec_locals = {}
        
        # 执行代码
        exec(code, exec_globals, exec_locals)
        print("  代码执行成功")
        
        # 获取指定目录下所有的图片文件
        image_files = tools.get_all_files_path('visualization')
        print(f"  找到{len(image_files)}个图像文件")
        
        if not image_files or len(image_files) == 0:
            print("  警告: 没有找到图像文件")
            gr.Warning("代码执行完成，但没有生成图像文件")
            
            # 创建警告图
            plt.figure(figsize=(6, 4))
            plt.text(0.5, 0.5, '没有生成图像文件', ha='center', va='center', fontsize=18)
            plt.axis('off')
            warning_file = os.path.join(vis_dir, 'no_image_warning.png')
            plt.savefig(warning_file)
            plt.close()
            
            return gr.Gallery(label='Visualization Warning', 
                              value=[warning_file], 
                              preview=True, 
                              visible=True, 
                              columns=1, 
                              rows=1, 
                              object_fit='contain', 
                              height='auto')
        
        # 只保留最新生成的5个图像，按修改时间排序
        sorted_files = sorted(image_files, key=os.path.getmtime, reverse=True)[:5]
        print(f"  选中了{len(sorted_files)}个最新的图像文件进行显示")
        
        gr.Info("可视化已生成")
        return gr.Gallery(label='Visualization Results', 
                          value=sorted_files, 
                          preview=True, 
                          visible=True, 
                          columns=2, 
                          rows=2, 
                          object_fit='contain', 
                          height='auto')
                          
    except Exception as e:
        print(f"  错误: 执行代码时发生异常: {e}")
        print(traceback.format_exc())
        gr.Error(f"代码执行错误: {str(e)}")
        
        # 创建错误提示图
        plt.figure(figsize=(6, 4))
        plt.text(0.5, 0.5, f'代码执行错误:\n{str(e)}', ha='center', va='center', fontsize=14)
        plt.axis('off')
        error_file = os.path.join('visualization', 'execution_error.png')
        plt.savefig(error_file)
        plt.close()
        
        return gr.Gallery(label='Execution Error', 
                          value=[error_file], 
                          preview=True, 
                          visible=True, 
                          columns=1, 
                          rows=1, 
                          object_fit='contain', 
                          height='auto')
